# python -m lilab.OpenLabCluster_train.a6_clippredpkl_2_seqencepkl clippredfile
#%%
import pickle
import numpy as np
import scipy.io as sio
import tqdm
import os.path as osp
import argparse
from lilab.openlabcluster_postprocess.s1a_clipNames_inplace_parse import parse_name


clippredfile = '/mnt/liying.cibr.ac.cn_Data_Temp/multiview_9/chenxf/00_BehaviorAnalysis-seq2seq/SexAge/Day55_Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/FWPCA0.00_P100_en3_hid30_epoch264-decSeqPC0.9_svm2allAcc0.94_kmeansK2use-42_fromK1-20_K100.clippredpkl'

realStart = 0
realEnd = 27000
fps=30
seg_length=int(fps*0.8)
gapTh=seg_length
fs=1  ##head and tail gap


def get_behavior_byFrame(bhvEachS,sfsEachS,realStart,realEnd):  ##bhvEachS corresponded to sfs
    ##get used video segment
    #all use relative start 0
    behavior=np.zeros(realEnd-realStart, dtype=int)
    sfsEachS=[i-realStart for i in sfsEachS]
    ##calculate transition by frame
    for si,s in enumerate(sfsEachS):behavior[s:s+seg_length]=bhvEachS[si]
    ##fill gaps
    for si in range(len(sfsEachS)-1):
       gap=sfsEachS[si+1]-sfsEachS[si]-seg_length
       if gap<gapTh and gap>0:
         behavior[sfsEachS[si]+seg_length:sfsEachS[si+1]]=bhvEachS[si]#fill by early behavior
    ##fill fs header and tail
    if sfsEachS[0]<=fs:   behavior[:fs]=bhvEachS[0]
    if sfsEachS[-1]+seg_length>=len(behavior)-fs:behavior[-fs:]=bhvEachS[-1]
    return behavior


def main(clippredfile:str, autoend:bool=False):
    clippreddata = pickle.load(open(clippredfile,'rb'))
    cluster_labels = clippreddata['cluster_labels'] #labels start from 1
    clipNames = clippreddata['clipNames']
    if 'df_clipNames' in clippreddata:
        df_clipNames = clippreddata['df_clipNames']
    else:
        df_clipNames = parse_name(clipNames)
    df_clipNames['cluster_label'] = cluster_labels
    bhv_seqFile = osp.splitext(clippredfile)[0]+'_sequences.pkl'

    bhv_dicts={}
    for (vnake, isBlack), df_clipIn in tqdm.tqdm(df_clipNames.groupby(['vnake', 'isBlack'])):
        df_clipIn = df_clipIn.sort_values('startFrame')
        fileName = f"fps{fps}_{vnake.replace('-','_')}_startFrame{realStart}_{'black' if isBlack else 'white'}First"
        sfsEachS, bhvEachS = df_clipIn.sort_values('startFrame')[['startFrame', 'cluster_label']].values.T
        the_end = sfsEachS.max()+seg_length if autoend else realEnd
        bhv= get_behavior_byFrame(bhvEachS,sfsEachS,realStart,the_end)
        bhv_dicts[fileName] = bhv

    sio.savemat(bhv_seqFile.replace('.pkl', '.mat'),bhv_dicts)
    pickle.dump(bhv_dicts,open(bhv_seqFile,'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('clippredfile', type=str, help='clippredfile')
    parser.add_argument('--autoEnd', action='store_true', help='auto check the last frame')
    args = parser.parse_args()
    main(args.clippredfile, args.autoEnd)
